#include "head.cuh"

void init_coo_mat(CooMat &cooUIview, CooMat &cooUIview_test, CooMat &cooUUview, Params params){

	cooUIview.row_di = cooUIview_test.row_di = nUsers;
	cooUIview.col_di = cooUIview_test.col_di = nItems;
	cooUUview.row_di = cooUUview.col_di = nUsers;

	cooUIview.nnz = (int)(params.split_rate*nRatings + 0.5) + 10000;
	cooUIview.rows_ind = (int*)malloc(cooUIview.nnz*sizeof(int));
	cooUIview.cols_ind = (int*)malloc(cooUIview.nnz*sizeof(int));
	cooUIview.vals = (float*)malloc(cooUIview.nnz*sizeof(float));

	cooUIview_test.nnz = nRatings - cooUIview.nnz + 10000;
	cooUIview_test.rows_ind = (int*)malloc(cooUIview_test.nnz*sizeof(int));
	cooUIview_test.cols_ind = (int*)malloc(cooUIview_test.nnz*sizeof(int));
	cooUIview_test.vals = (float*)malloc(cooUIview_test.nnz*sizeof(float));

	cooUUview.nnz = nLinks * 2 + nUsers + 10000;
	cooUUview.rows_ind = (int*)malloc(cooUUview.nnz*sizeof(int));
	cooUUview.cols_ind = (int*)malloc(cooUUview.nnz*sizeof(int));
	cooUUview.vals = (float*)malloc(cooUUview.nnz*sizeof(float));
}

void init_csr_mat_device(CsrMatOnDevice &d_csrUIview, CooMatOnDevice &d_cooUIview, CsrMatOnDevice &d_csrUIview_test,  CooMatOnDevice &d_cooUIview_test, CsrMatOnDevice &d_csrIUview, CsrMatOnDevice &d_csrUUview, CooMatOnDevice &d_cooUUview){
	d_csrIUview.col_di = d_csrUIview.row_di = d_csrUIview_test.row_di = d_cooUIview.row_di;
	d_csrIUview.row_di = d_csrUIview.col_di = d_csrUIview_test.col_di = d_cooUIview.col_di;
	d_csrUUview.row_di = d_csrUUview.col_di = d_cooUIview.col_di;

	d_csrIUview.nnz = d_csrUIview.nnz = d_cooUIview.nnz;
	d_csrUIview_test.nnz = d_cooUIview_test.nnz;
	d_csrUUview.nnz = d_cooUUview.nnz;

	CHECK_CUDA(cudaMalloc(&(d_csrUIview.vals), d_csrUIview.nnz*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&(d_csrUIview.row_ptr), (d_csrUIview.row_di+1)*sizeof(int)));
	CHECK_CUDA(cudaMalloc(&(d_csrUIview.cols_ind), d_csrUIview.nnz*sizeof(int)));

	CHECK_CUDA(cudaMalloc(&(d_csrIUview.vals), d_csrIUview.nnz*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&(d_csrIUview.row_ptr), (d_csrIUview.row_di+1)*sizeof(int)));
	CHECK_CUDA(cudaMalloc(&(d_csrIUview.cols_ind), d_csrIUview.nnz*sizeof(int)));

	d_csrUIview_test.vals = d_cooUIview_test.vals;
	CHECK_CUDA(cudaMalloc(&(d_csrUIview_test.row_ptr), (d_csrUIview_test.row_di+1)*sizeof(int)));
	d_csrUIview_test.cols_ind = d_cooUIview_test.cols_ind;

	d_csrUUview.vals = d_cooUUview.vals;
	CHECK_CUDA(cudaMalloc(&(d_csrUUview.row_ptr), (d_csrUUview.row_di+1)*sizeof(int)));
	d_csrUUview.cols_ind = d_cooUUview.cols_ind;
}

void copy_coo_mat_host_2_device(CooMat &hcm, CooMatOnDevice &dcm){
	dcm.nnz = hcm.nnz;
	dcm.row_di = hcm.row_di;
	dcm.col_di = hcm.col_di;
	CHECK_CUDA(cudaMalloc(&(dcm.rows_ind),dcm.nnz*sizeof(int)));
	CHECK_CUDA(cudaMemcpy(dcm.rows_ind, hcm.rows_ind, dcm.nnz*sizeof(int), cudaMemcpyHostToDevice));
	CHECK_CUDA(cudaMalloc(&(dcm.cols_ind), dcm.nnz*sizeof(int)));
	CHECK_CUDA(cudaMemcpy(dcm.cols_ind, hcm.cols_ind, dcm.nnz*sizeof(int), cudaMemcpyHostToDevice));
	CHECK_CUDA(cudaMalloc(&(dcm.vals), dcm.nnz*sizeof(int)));
	CHECK_CUDA(cudaMemcpy(dcm.vals, hcm.vals, dcm.nnz*sizeof(float), cudaMemcpyHostToDevice));
}

void free_coo_mat_host(CooMat &cm){
	free(cm.cols_ind); cm.cols_ind = NULL;
	free(cm.rows_ind); cm.rows_ind = NULL;
	free(cm.vals); cm.vals = NULL;
}

void coo_mat_sort_by_row(CooMatOnDevice &d_coo_mat, cusparseMatDescr_t descr, cusparseHandle_t handle){

	size_t pBufferSizeInBytes = 0;
	void *pBuffer = NULL;
	int *P = NULL;
	CHECK_CUSPARSE(cusparseXcoosort_bufferSizeExt(handle, d_coo_mat.row_di, d_coo_mat.col_di, d_coo_mat.nnz, 
		d_coo_mat.rows_ind, d_coo_mat.cols_ind, &pBufferSizeInBytes));
	CHECK_CUDA(cudaMalloc(&pBuffer, sizeof(char)* pBufferSizeInBytes));
	CHECK_CUDA(cudaMalloc((void**)&P, sizeof(int)*d_coo_mat.nnz));
	CHECK_CUSPARSE(cusparseCreateIdentityPermutation(handle, d_coo_mat.nnz, P));
	CHECK_CUSPARSE(cusparseXcoosortByRow(handle, d_coo_mat.row_di, d_coo_mat.col_di, d_coo_mat.nnz,
		d_coo_mat.rows_ind, d_coo_mat.cols_ind, P, pBuffer));
	float *cooVals_sorted;
	CHECK_CUDA(cudaMalloc(&cooVals_sorted, d_coo_mat.nnz*sizeof(float)));
	CHECK_CUSPARSE(cusparseSgthr(handle, d_coo_mat.nnz, d_coo_mat.vals, cooVals_sorted, P, CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaMemcpy(d_coo_mat.vals, cooVals_sorted, d_coo_mat.nnz*sizeof(float), cudaMemcpyDeviceToDevice));
	CHECK_CUDA(cudaFree(cooVals_sorted));
	CHECK_CUDA(cudaFree(P));
	CHECK_CUDA(cudaFree(pBuffer));
}

void coo_mat_sort_by_col(CooMatOnDevice &d_coo_mat, cusparseMatDescr_t descr, cusparseHandle_t handle){

	size_t pBufferSizeInBytes = 0;
	void *pBuffer = NULL;
	int *P = NULL;
	CHECK_CUSPARSE(cusparseXcoosort_bufferSizeExt(handle, d_coo_mat.row_di, d_coo_mat.col_di, d_coo_mat.nnz,
		d_coo_mat.rows_ind, d_coo_mat.cols_ind, &pBufferSizeInBytes));
	CHECK_CUDA(cudaMalloc(&pBuffer, sizeof(char)* pBufferSizeInBytes));
	CHECK_CUDA(cudaMalloc((void**)&P, sizeof(int)*d_coo_mat.nnz));
	CHECK_CUSPARSE(cusparseCreateIdentityPermutation(handle, d_coo_mat.nnz, P));
	CHECK_CUSPARSE(cusparseXcoosortByColumn(handle, d_coo_mat.row_di, d_coo_mat.col_di, d_coo_mat.nnz,
		d_coo_mat.rows_ind, d_coo_mat.cols_ind, P, pBuffer));
	float *cooVals_sorted;
	CHECK_CUDA(cudaMalloc(&cooVals_sorted, d_coo_mat.nnz*sizeof(float)));
	CHECK_CUSPARSE(cusparseSgthr(handle, d_coo_mat.nnz, d_coo_mat.vals, cooVals_sorted, P, CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaMemcpy(d_coo_mat.vals, cooVals_sorted, d_coo_mat.nnz*sizeof(float), cudaMemcpyDeviceToDevice));
	CHECK_CUDA(cudaFree(cooVals_sorted));
	CHECK_CUDA(cudaFree(P));
	CHECK_CUDA(cudaFree(pBuffer));
}

void copy_to_host(CsrMat &h_csr_mat, CsrMatOnDevice &d_csr_mat){
	h_csr_mat.nnz = d_csr_mat.nnz;
	h_csr_mat.row_di = d_csr_mat.row_di;
	h_csr_mat.col_di = d_csr_mat.col_di;
	h_csr_mat.vals = (float*)malloc(h_csr_mat.nnz*sizeof(float));
	CHECK_CUDA(cudaMemcpy(h_csr_mat.vals, d_csr_mat.vals, h_csr_mat.nnz*sizeof(float),cudaMemcpyDeviceToHost));
	h_csr_mat.row_ptr = (int*)malloc((h_csr_mat.row_di + 1)*sizeof(int));
	CHECK_CUDA(cudaMemcpy(h_csr_mat.row_ptr, d_csr_mat.row_ptr,(h_csr_mat.row_di+1)*sizeof(int),cudaMemcpyDeviceToHost));
	h_csr_mat.cols_ind = (int*)malloc(h_csr_mat.nnz*sizeof(int));
	CHECK_CUDA(cudaMemcpy(h_csr_mat.cols_ind,d_csr_mat.cols_ind,h_csr_mat.nnz*sizeof(int),cudaMemcpyDeviceToHost));
}

void ffff(float *h_anchors, int d, int n, char *filename){
	FILE *f;
	f = fopen(filename, "w");
	std::stringstream ss;
	for (int i = 0; i < d; i++){
		ss.str("");
		for (int j = 0; j < n; j++){
			ss << h_anchors[i+ d*j] << " ";
		}
		ss << "\n";
		fprintf(f, "%s", ss.str().c_str());
	}
	fclose(f);
}

int main(void){
	Params params;
	params.split_rate = 0.8;
	params.r = 8;
	params.maxS = 5;
	params.minS = 1;
	params.num_of_kmeans = 3;
	params.num_of_LAE_iters = 10;
	params.num_of_anchors = 20;
	params.num_of_nearest_anchors = 3;
	params.alpha = 1;
	params.beta = 1;
	params.rho1 = 1;
	params.rho2 = 1;

	/* 
	Step 1
	read dataset into 4 CsrMatOnDevice :
		d_csrUIview: nUsers X nItems
		d_csrUUview: nUsers X nUsers
		d_csrUIview_test: nUsers X nItems
		d_csrIUview: nItems X nUsers
	*/
	clock_t read_data_start = clock();
	CooMat cooUIview, cooUIview_test, cooUUview;
	init_coo_mat(cooUIview, cooUIview_test, cooUUview, params);
	read_data_2_coo_mat(cooUIview, cooUIview_test, cooUUview, params.split_rate);

	CooMatOnDevice d_cooUIview, d_cooUIview_test, d_cooUUview;
	copy_coo_mat_host_2_device(cooUIview, d_cooUIview);
	copy_coo_mat_host_2_device(cooUIview_test, d_cooUIview_test);
	copy_coo_mat_host_2_device(cooUUview, d_cooUUview);
	free_coo_mat_host(cooUIview);
	free_coo_mat_host(cooUIview_test);
	free_coo_mat_host(cooUUview);

	cusparseMatDescr_t descr;
	cusparseCreateMatDescr(&descr);
	CHECK_CUSPARSE(cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL));
	CHECK_CUSPARSE(cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO));
	cusparseHandle_t handle;
	CHECK_CUSPARSE(cusparseCreate(&handle));

	CsrMatOnDevice d_csrUIview, d_csrUUview, d_csrIUview, d_csrUIview_test;
	init_csr_mat_device(d_csrUIview, d_cooUIview, 
		d_csrUIview_test, d_cooUIview_test, 
		d_csrIUview,
		d_csrUUview, d_cooUUview);

	coo_mat_sort_by_row(d_cooUIview, descr, handle);
	CHECK_CUSPARSE(cusparseXcoo2csr(handle, d_cooUIview.rows_ind, d_cooUIview.nnz, d_cooUIview.row_di, d_csrUIview.row_ptr,
		CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaMemcpy(d_csrUIview.cols_ind, d_cooUIview.cols_ind, d_cooUIview.nnz*sizeof(int), cudaMemcpyDeviceToDevice));
	CHECK_CUDA(cudaMemcpy(d_csrUIview.vals, d_cooUIview.vals, d_cooUIview.nnz*sizeof(float), cudaMemcpyDeviceToDevice));

	coo_mat_sort_by_col(d_cooUIview, descr, handle);
	CHECK_CUSPARSE(cusparseXcoo2csr(handle, d_cooUIview.cols_ind, d_cooUIview.nnz, d_cooUIview.col_di, d_csrIUview.row_ptr,
		CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaMemcpy(d_csrIUview.cols_ind, d_cooUIview.rows_ind, d_cooUIview.nnz*sizeof(int), cudaMemcpyDeviceToDevice));
	CHECK_CUDA(cudaMemcpy(d_csrIUview.vals, d_cooUIview.vals, d_cooUIview.nnz*sizeof(float), cudaMemcpyDeviceToDevice));

	CHECK_CUDA(cudaFree(d_cooUIview.cols_ind)); CHECK_CUDA(cudaFree(d_cooUIview.rows_ind)); CHECK_CUDA(cudaFree(d_cooUIview.vals));

	coo_mat_sort_by_row(d_cooUIview_test, descr, handle);
	CHECK_CUSPARSE(cusparseXcoo2csr(handle, d_cooUIview_test.rows_ind, d_cooUIview_test.nnz, d_cooUIview_test.row_di, d_csrUIview_test.row_ptr,
		CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaFree(d_cooUIview_test.rows_ind));

	coo_mat_sort_by_row(d_cooUUview, descr, handle);
	CHECK_CUSPARSE(cusparseXcoo2csr(handle, d_cooUUview.rows_ind, d_cooUUview.nnz, d_cooUUview.row_di, d_csrUUview.row_ptr,
		CUSPARSE_INDEX_BASE_ZERO));
	CHECK_CUDA(cudaFree(d_cooUUview.rows_ind));

	CsrMat h_csrUIview, h_csrUUview, h_csrIUview;
	copy_to_host(h_csrUIview, d_csrUIview);
	copy_to_host(h_csrIUview, d_csrIUview);
	copy_to_host(h_csrUUview, d_csrUUview);

	CHECK_CUDA(cudaDeviceSynchronize());
	CHECK_CUDA(cudaGetLastError());
	
	clock_t read_data_finish = clock();
	printf("read dataset done. cost time: %f\n", (double)(read_data_finish - read_data_start) / CLOCKS_PER_SEC);

	/*
	Step 2
	get anchors of 3 CsrMatOnDevice: d_csrUIview, d_csrIUview, d_csrUUview
	each row of CsrMatOnDevice stands for its feature vectors
	*/
	float *d_anchors_of_uiview, *d_anchors_of_iuview, *d_anchors_of_uuview;
	// Note that anchors are dense matrix: d X n
	CHECK_CUDA(cudaMalloc(&d_anchors_of_uiview, d_csrUIview.col_di*params.num_of_anchors*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_anchors_of_iuview, d_csrIUview.col_di*params.num_of_anchors*sizeof(float)));
	CHECK_CUDA(cudaMalloc(&d_anchors_of_uuview, d_csrUUview.col_di*params.num_of_anchors*sizeof(float)));

	get_anchors_by_Kmeans(d_anchors_of_uiview, d_csrUIview, params);
	clock_t t1 = clock();
	printf("get UIview anchors done. cost time: %f\n", (double)(t1 - read_data_finish) / CLOCKS_PER_SEC);

	get_anchors_by_Kmeans(d_anchors_of_iuview, d_csrIUview, params);
	clock_t t2 = clock();
	printf("get IUview anchors done. cost time: %f\n", (double)(t2 - t1) / CLOCKS_PER_SEC);

	get_anchors_by_Kmeans(d_anchors_of_uuview, d_csrUUview, params);
	clock_t t3 = clock();
	printf("get UUview anchors done. cost time: %f\n", (double)(t3 - t2) / CLOCKS_PER_SEC);

	CHECK_CUSPARSE(cusparseDestroyMatDescr(descr));
	CHECK_CUSPARSE(cusparseDestroy(handle));
	
	float *h_anchors_of_uiview = (float*)malloc(d_csrUIview.col_di*params.num_of_anchors*sizeof(float));
	float *h_anchors_of_iuview = (float*)malloc(d_csrIUview.col_di*params.num_of_anchors*sizeof(float));
	float *h_anchors_of_uuview = (float*)malloc(d_csrUUview.col_di*params.num_of_anchors*sizeof(float));
	CHECK_CUDA(cudaMemcpy(h_anchors_of_uiview, d_anchors_of_uiview, d_csrUIview.col_di*params.num_of_anchors*sizeof(float),
		cudaMemcpyDeviceToHost));
	CHECK_CUDA(cudaMemcpy(h_anchors_of_iuview, d_anchors_of_iuview, d_csrIUview.col_di*params.num_of_anchors*sizeof(float),
		cudaMemcpyDeviceToHost));
	CHECK_CUDA(cudaMemcpy(h_anchors_of_uuview, d_anchors_of_uuview, d_csrUUview.col_di*params.num_of_anchors*sizeof(float),
		cudaMemcpyDeviceToHost));

	cublasHandle_t cublashandle;
	CHECK_CUDA_STATUS(cublasCreate(&cublashandle));
	
	float *d_uZ, *h_uZ;
	h_uZ = (float*)malloc(nUsers*(2 * params.num_of_anchors)*sizeof(float));
	memset(h_uZ, 0, nUsers*(2 * params.num_of_anchors)*sizeof(float));
	get_Z(h_uZ, params.num_of_anchors, d_csrUIview, h_csrUIview, d_anchors_of_uiview, h_anchors_of_uiview, d_csrUUview, h_csrUUview, d_anchors_of_uuview, h_anchors_of_uuview, params, cublashandle);
	CHECK_CUDA(cudaMalloc(&d_uZ, nUsers*(2 * params.num_of_anchors)*sizeof(float)));
	CHECK_CUDA(cudaMemcpy(d_uZ, h_uZ, nUsers*(2 * params.num_of_anchors)*sizeof(float), cudaMemcpyHostToDevice));

	float *d_iZ, *h_iZ;
	h_iZ = (float*)malloc(nItems*params.num_of_anchors*sizeof(float));
	memset(h_iZ, 0, nItems*params.num_of_anchors*sizeof(float));
	get_Z(h_iZ, params.num_of_anchors, d_csrIUview, h_csrIUview, d_anchors_of_iuview, h_anchors_of_iuview, params, cublashandle);
	CHECK_CUDA(cudaMalloc(&d_iZ, nItems*params.num_of_anchors*sizeof(float)));
	CHECK_CUDA(cudaMemcpy(d_iZ, h_iZ, nItems*params.num_of_anchors*sizeof(float), cudaMemcpyHostToDevice));

	//free(h_anchors_of_iuview); free(h_anchors_of_uiview); free(h_anchors_of_uuview);


	CHECK_CUDA_STATUS(cublasDestroy(cublashandle));
//
//	float *d_S_u;
//	CHECK_CUDA(cudaMalloc(&d_S_u, nUsers*nUsers*sizeof(float)));
//	float *d_S_i;
//	CHECK_CUDA(cudaMalloc(&d_S_i,nItems*nItems*sizeof(float)));
//	get_S_from_Z(d_S_u, uZ, nUsers, 2 * params.num_of_anchors, cublashandle);
//	get_S_from_Z(d_S_i, iZ, nItems, params.num_of_anchors, cublashandle);
//
//	free(iZ);
//	free(uZ);
//
//	// now we have S_i and S_u
//	// next step is to calculate B and D by MVDCF
//	// B is r-bits code of users: r X nusers
//	// D is r-bits code of items: r X nitems
//	float *B = (float *)malloc(params.r*nUsers*sizeof(float));
//	float *D = (float *)malloc(params.r*nItems*sizeof(float));
//
//	MVDCF(B, D, params, dev_UIview, dev_IUview, d_S_i, d_S_u, cublashandle);
//	write_into_file(B, params.r, nUsers, "B_code.txt");
//	write_into_file(D, params.r, nItems, "D_code.txt");
//	// save UIview_test
//	write_test_data(UIview_test, nUsers, nItems, "UIview_test.txt");
//	printf("write B and D into files done.\n");
//
//	//int result_num = 5;
//	//evaluate(B, D, params.r, nUsers, nItems, UIview_test,
//	//	result_num, cublashandle);
//
//	CHECK_CUDA_STATUS(cublasDestroy(cublashandle));
//	goto CLEAN_WORK;
//CLEAN_WORK:
//	CHECK_CUDA(cudaFree(dev_UIview));
//	CHECK_CUDA(cudaFree(dev_IUview));
//	CHECK_CUDA(cudaFree(d_S_i));
//	CHECK_CUDA(cudaFree(d_S_u));
//	free(UIview);
//	free(UIview_test);
//	free(B);
//	free(D);
//	CHECK_CUDA(cudaDeviceReset());

	system("pause");
	return EXIT_SUCCESS; 
}